package com.feiwo.nonego.rpc.registry;

import com.feiwo.nonego.rpc.common.util.RpcServiceUtils;
import com.feiwo.nonego.rpc.loadbalancer.api.ServiceLoadBalancer;
import com.feiwo.nonego.rpc.protocol.meta.ServiceMeta;
import com.feiwo.nonego.rpc.registry.api.RegistryService;
import com.feiwo.nonego.rpc.registry.api.config.RegistryConfig;
import com.feiwo.nonego.rpc.spi.loader.ExtensionLoader;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;

import java.io.IOException;
import java.util.List;
import java.util.Objects;

/**
 *
 *
 * @author feiwo 2025-08-08
 */
public class ZookeeperRegistryService implements RegistryService {

    private static final int BASE_SLEEP_TIME_MS = 30000;

    private static final int MAX_RETRIES = 3;

    private static final String ZK_BASE_PATH = "/nonego_rpc";

    private ServiceDiscovery<ServiceMeta> serviceDiscovery;
    private ServiceLoadBalancer<ServiceInstance<ServiceMeta>> serviceLoadBalancer;

    @Override
    public void init(RegistryConfig registryConfig) throws Exception {
        var client = CuratorFrameworkFactory.newClient(registryConfig.getRegistryAddress(), new ExponentialBackoffRetry(BASE_SLEEP_TIME_MS, MAX_RETRIES));
        client.start();

        var serializer = new JsonInstanceSerializer<>(ServiceMeta.class);
        this.serviceDiscovery = ServiceDiscoveryBuilder.builder(ServiceMeta.class)
                .client(client)
                .serializer(serializer)
                .basePath(ZK_BASE_PATH)
                .build();
        this.serviceDiscovery.start();
        this.serviceLoadBalancer = ExtensionLoader.getExtension(ServiceLoadBalancer.class, registryConfig.getLoadBalancerType());
    }

    @Override
    public void register(ServiceMeta meta) throws Exception {
        ServiceInstance<ServiceMeta> serviceInstance = ServiceInstance
                .<ServiceMeta>builder()
                .name(RpcServiceUtils.buildServiceKey(meta.getServiceName(), meta.getServiceVersion(), meta.getServiceGroup()))
                .address(meta.getServiceAddress())
                .port(meta.getServicePort())
                .payload(meta)
                .build();
        serviceDiscovery.registerService(serviceInstance);
    }

    @Override
    public void unRegister(ServiceMeta meta) throws Exception {
        ServiceInstance<ServiceMeta> serviceInstance = ServiceInstance
                .<ServiceMeta>builder()
                .name(RpcServiceUtils.buildServiceKey(meta.getServiceName(), meta.getServiceVersion(), meta.getServiceGroup()))
                .address(meta.getServiceAddress())
                .port(meta.getServicePort())
                .payload(meta)
                .build();
        serviceDiscovery.unregisterService(serviceInstance);
    }

    @Override
    public ServiceMeta discovery(String serviceName, int invokerHashCode, String sourceIp) throws Exception {
        var serviceInstances = serviceDiscovery.queryForInstances(serviceName);
        var instance = serviceLoadBalancer.select((List<ServiceInstance<ServiceMeta>>) serviceInstances, invokerHashCode, sourceIp);
        if (Objects.nonNull(instance)) {
            return instance.getPayload();
        }
        return null;
    }

    @Override
    public void destroy() throws IOException {
        serviceDiscovery.close();
    }
}
